"# -*- encoding: utf-8 -*-"
"""
@Project :   quant
@File    :   data_utils.py
@Time    :   2019/11/15 13:21
@Author  :   BiKang Peng
@Software:   PyCharm
@Contact :   782222720@qq.com
@License :   (C)Copyright 2017-2018, Liugroup-NLPR-CASIA
@Desc : ==============================================
Life is Short I Use Python!!!                      ===
If this runs wrong,don't ask me,I don't know why;  ===
If this runs right,thank god,and I don't know why. ===
Maybe the answer,my friend,is blowing in the wind. ===
======================================================
"""
import tushare as ts
import pandas as pd
import os
import time
import glob
from stock import *

ts.set_token('52d1d7b3ade4d917050699b12c7f2327fd9c2afa2e21e66f2633a88b')
pro = ts.pro_api()

def getParameter(ts_code=None):
    # 切割起始时间，按年获取数据，再进行拼接
    # 返回一个parameters的列表
    start = int(STARTDATE / 10000)
    end = int(ENDDATE / 10000)
    paralist = []
    for i in range(start, end + 1):
        para = Parameters(ts_code=ts_code,
                          start_date=str(i) + '0101',
                          end_date=str(i) + '1231')
        paralist.append(para)
    return paralist

def getDailyStock(paralist):
    '''
    获取每日的股价数据以及基本数据
    一边获取数据 一边修改列名
    '''
    total = pd.DataFrame()
    for para in paralist:
        stockdata = StockData(para)
        cal = stockdata.getTradeCalender().drop(columns=['exchange', 'is_open'])
        daily = stockdata.getDaily().drop(columns='ts_code')
        df = pd.merge(cal, daily, left_on='cal_date', right_on='trade_date', how='left')
        total = total.append(df.sort_values(by='cal_date', ascending=True), ignore_index=True)
    print('Get {0} stock market data at {1} dimentions and {2} rows.'.format(TSCODE, total.shape[1], total.shape[0]))
    total.to_csv('./stocks_data/stocks_training/' + TSCODE + '.csv')
    return total

# ----------------------下载某只股票数据------------------- #
# def get_stock_data(code, start_date, end_date):
#     '''
#     获取A股2014年前上市的股票数据
#     :param code: 股票代码
#     :param start_date: 开始日期
#     :param end_date: 结束日期
#     :param filename: 保存的目录
#     :param length:
#     :return:
#     '''
#     df = pro.daily(ts_code=code, start_date=start_date, end_date=end_date)
#     df = df[['trade_date', 'open', 'high', 'low', 'close', 'vol', 'amount', 'pct_chg']]
#     # 排序，日期从20191104-20110104 变成20110104-20191104
#     df = df.sort_values(by='trade_date', ascending=True)
#     # 重新排序索引
#     df = df.reset_index(drop=True)
#     # df = df.drop(columns=['Unnamed: 0'])
#     print('共有%s天数据' % len(df))
#     # if length == -1:
#     #     path = code + '.csv'
#     #     df.to_csv(os.path.join(filename, path))
#     # else:
#     #     if len(df1) >= length:
#     #         path = code + '.csv'
#     #         df.to_csv(os.path.join(filename, path))


# ----------------------下载沪深300指数数据------------------- #
def get_hs300_data(start_date, end_date, filname):
    '''
    获取沪深300股指数据
    :param start_date:
    :param end_dat:
    :param filname:
    :return:
    '''
    df = pro.daily('399300.SZ', start_date=start_date, end_date=end_date)
    df = df[['trade_date', 'open', 'high', 'low', 'close', 'vol', 'amount', 'pct_chg']]
    df = df.sort_values(by='trade_date', ascending=True)
    df = df.reset_index(drop=True)
    print('共有%s天数据' % len(df))
    df.to_csv(os.path.join(filename, '399300.csv'))


# ------------------------更新股票数据------------------------ #
def update_stock_data(filename):
    '''
    更新股票数据
    :param filename:
    :return:
    '''
    (filepath, tempfilename) = os.path.split(filename)
    (stock_code, extension) = os.path.splitext(tempfilename)
    f = open(filename, 'r')
    df = pd.read_csv(f)
    print('股票{}文件中的最新日期为:{}'.format(stock_code, df.iloc[-1, 0]))
    data_now = time.strftime('%Y-%m-%d', time.localtime(time.time()))
    print('更新日期至：%s' % data_now)
    nf = pro.daily(stock_code, str(df.iloc[-1, 0]), data_now)
    nf = nf.sort_values(by='trade_date', ascending=True)
    nf = nf.reset_index(drop=True)
    nf = nf.iloc[1:]
    print('共有%s天数据' % len(nf))
    nf = nf[['trade_date', 'open', 'high', 'low', 'close', 'vol', 'amount', 'pct_chg']]
    nf.to_csv(filename, mode='a', header=False)
    f.close()


# ------------------------获取股票长度----------------------- #
def get_data_len(file_path):
    with open(file_path) as f:
        df = pd.read_csv(f)
        return len(df)


# --------------------------日期筛选------------------------- #
# 对已经再本地的文件按照日期筛选，date1是开始数据，date2是结束数据
# file1是源文件夹，file2是筛选日期后文件存放的文件夹
def select_stock_data(file1, file2, state_date, end_date):
    '''

    :param file1: 源文件夹
    :param file2: 筛选日期后文件存放的文件夹
    :param state_date: 开始数据
    :param end_date: 结束数据
    :return:
    '''
    csv_list = glob.glob(file1 + '*.csv')
    print(u'共发现%s个CSV文件' % len(csv_list))
    file_list = []
    for i in csv_list:
        (filepath, filename) = os.path.split(i)
        file_list.append(filename)
    for i in file_list:
        f = open(os.path.join(file1, i), 'r')
        df1 = pd.read_csv(f, header=0)
        df1['date'] = pd.to_datetime(df1['date'])
        df1 = df1.set_index('date')
        df2 = df1[state_date:end_date]
        df2.to_csv(os.path.join(file2, i))

# --------------------------停盘填充------------------------- #
def fill_stock_data(target, sfile, tfile):
    '''
    按照沪深300指数来对个股停盘数据进行填充，填充为该股上一交易日的数据
    :param target: 参照股票
    :param sfile: 原文件夹
    :param tfile: 填充完要存放文件夹
    :return:
    '''
    tf = open(target)
    tf = pd.read_csv(tf)
    csv_list = glob.glob(sfile + '*.csv')
    print(u'共发现%s个CSV文件' % len(csv_list))
    i = 1
    for item in csv_list:
        f1 = open(item)
        print('正在处理第%s个文件' % i)
        df2 = pd.read_csv(f1)
        mix_data = pd.merge(tf, df2, how='outer', on="trade_date")
        mix_data = mix_data.fillna(method='pad')
        d1 = mix_data[['trade_date', 'open_y', 'high_y', 'low_y', 'close_y', 'vol_y', 'amount_y', 'pct_chg_y']]
        d1.rename(columns={'open_y': 'open', 'high_y': 'high', 'low_y': 'low', 'close_y': 'close', 'vol_y': 'vol', 'amount_y': 'amount', 'pct_chg_y': 'pct_chg'}, inplace=True)
        (filepath, filename) = os.path.split(item)
        d1.to_csv(os.path.join(tfile, filename), index=False)
        i += 1

# --------------------------文件合并------------------------- #
def merge_stock_data(filename, tfile):
    '''
    将多个文件合并为一个文件，在文件末尾添加
    :param filename: 需要合并的文件夹
    :param tfile: 存放合并后文件的文件夹
    :return:
    '''
    csv_list = glob.glob(filename + '*.csv')
    print(u'共发现%s个CSV文件' % len(csv_list))
    f = open(csv_list[0])
    df = pd.read_csv(f)
    for i in range(1, len(csv_list)):
        f1 = open(csv_list[i], 'rb')
        df1 = pd.read_csv(f1)
        df = pd.concat([df, df1])
    df.to_csv(tfile+'train_mix.csv', index=None)

if __name__ == '__main__':
    STARTDATE = 20150101
    ENDDATE = 20171231
    codelist = []
    with open('stocklist.csv', encoding="UTF-8") as f:
        df = pd.read_csv(f, converters={'ts_code': str})
        codelist.extend(df['ts_code'])
    i = 1
    for code in codelist:
        TSCODE = code
        paralist = getParameter(ts_code=TSCODE)
        df = getDailyStock(paralist)
        time.sleep(5)


